import os
import uuid
import cv2
import numpy as np
import subprocess
import shutil
import multiprocessing as mp
import pandas as pd
from pandas import DataFrame as df
from math import log,sqrt
from sklearn import preprocessing
from math import sin,cos,pi,pow,log10
from scipy.interpolate import interp1d
from scipy.misc import derivative

class KNOW:
    def __init__(self,nlayer,first,zrate,zexpand,xrange,interval,xrate,xexpand):
        self.NLAYER=nlayer
        self.FIRST=first
        self.ZRATE=zrate
        self.ZEXPAND=zexpand
        self.XRANGE=xrange
        self.INTERVAL=interval
        self.XRATE=xrate
        self.XEXPAND=xexpand
        self.COOR=[]
        self.tic()
        self.t2d()
        self.recv()
    
    def tic(self):
        self.TIC,temp=[self.FIRST],self.FIRST
        for i in range(self.NLAYER-1):
            temp=temp*self.ZRATE
            self.TIC.append(temp)
        temp=[det*self.FIRST for det in [i*self.XRATE for i in range(1,self.ZEXPAND+1)]]
        [self.TIC.append(det) for det in temp]
        self.NLAYER=self.NLAYER+self.ZEXPAND
    
    def t2d(self):
        self.DEEP=[self.TIC[0]]
        for i in range(1,self.NLAYER):
            self.DEEP.append(self.TIC[i]+self.DEEP[i-1])
    
    def recv(self):
        array=[str(self.INTERVAL) for i in range(self.XRANGE)]
        temp=[str(det*self.INTERVAL) for det in [i*self.XRATE for i in range(1,int(self.XEXPAND)+1)]]
        [array.insert(0,det) for det in temp]
        [array.append(det) for det in temp]
        self.XRANGE=int(2*self.XEXPAND+self.XRANGE)
        self.INTERVAL=array

        array=[float(det) for det in array]
        temp,x0,coor=0.0,0.0,[]
        for i,_ in enumerate(array):
            temp=temp+(0.0 if i==0 else array[i-1])
            swap=x0+temp+array[i]/2.0
            coor.append(swap)
        self.COOR=[coor[int(self.XEXPAND+i)] for i in range(int(self.XRANGE-2*self.XEXPAND))]

class MEASURE:
    def __init__(self):
        self.NFREQ=len(self.GetFreqs())
    
    def Measured(self,filename,coor):
        freq=self.GetFreqs()
        with open(filename,"w") as fp:            
            fp.write("# 8b851557\n")
            fp.write("# Period(s) Code GG_Lat GG_Lon X(m) Y(m) Z(m) Component Value Error\n")
            fp.write("> TM_Impedance\n> exp(-i\omega t)\n> [V/m]/[A/m]\n> 0.00\n> 0.000 0.000\n")
            # fp.write("> Rho_Phs_TM\n> exp(-i\omega t)\n> []\n> 0.00\n> 0.000 0.000\n")
            fp.write("> "+" ".join([str(det) for det in [self.NFREQ,len(coor)]])+"\n")  
            for i,each in enumerate(coor):
                fp.write("".join([" ".join([str(row) for row in [det,"L0S"+str(i),0,0,0,each,0,"TM",0,0,0,"\n"]]) for det in freq]))
                # fp.write("".join([" ".join([str(row) for row in [det,"L0S"+str(i),0,0,0,each,0,"Rho_TM",0,0,"\n"]]) for det in freq]))
            fp.write("# 8b851557\n")
            fp.write("# Period(s) Code GG_Lat GG_Lon X(m) Y(m) Z(m) Component Value Error\n")
            fp.write("> TE_Impedance\n> exp(-i\omega t)\n> [V/m]/[A/m]\n> 0.00\n> 0.000 0.000\n")
            fp.write("> "+" ".join([str(det) for det in [self.NFREQ,len(coor)]])+"\n")  
            for i,each in enumerate(coor):
                fp.write("".join([" ".join([str(row) for row in [det,"L0S"+str(i),0,0,0,each,0,"TE",0,0,0,"\n"]]) for det in freq]))

        self.measuredOfFilename=filename

    def Forward(self,know,filename,result,target,zoom,nx,ny,xrange,nlayer,interval,thick,ground,index):
        image=know.reshape((target,target))
        top_border,bottom_border,left_border,right_border=int((target-ny*zoom)/2.),int((target-ny*zoom)/2.),int((target-nx*zoom)/2.),int((target-nx*zoom)/2.)
        image=image[top_border:target-bottom_border,left_border:target-right_border]
        image=np.log(((1.0-cv2.resize(image,(nx,ny)))*ground))
        image=np.nan_to_num(image,nan=1)

        i=str(index)

        with open(filename+i,"w") as fp:        
            fp.write(" ".join([str(det) for det in [xrange,nlayer,"LOGE"]])+"\n")
            fp.write(" ".join([str(det) for det in interval])+"\n")
            fp.write(" ".join([str(det) for det in thick])+"\n")
            np.savetxt(fp,image)
        
        subprocess.run(["mt2di"+i+".exe","-F",filename+i,self.measuredOfFilename,result+i],capture_output=False)     
        
        content=[]
        with open(result+i,"r") as fp:
            content.extend([det.split() for det in fp])    
        # content=content[8:]
        content=[det for det in content if len(det)==11]

        # return np.array([abs(complex(float(det[-3]),float(det[-2])))/3.0 for _,det in enumerate(content)]) 
        data=[]
        [data.extend([abs(float(det[-3])/3.0),abs(float(det[-2])/3.0)]) for _,det in enumerate(content)]
        return np.array(data)
        # return np.array([0.1*log10(float(det[-2]))+0.2 for _,det in enumerate(content)])   
    
    def GetFreqs(self):
        freq="0.03125 0.04419417382415922 0.0625 0.08838834764831845 0.125 0.1767766952966369 0.25 0.3535533905932738 0.5 0.7071067811865476 1.0 1.4142135623730951 2.0 2.8284271247461903 4.0 5.656854249492381 8.0 11.313708498984761 16.0 22.627416997969522 32.0 45.254833995939045 64.0 90.50966799187809 128.0 181.01933598375618 256.0 362.03867196751236 512.0 724.0773439350247 1024.0 1448.1546878700494 2048.0 2896.309375740099 4096.0 5792.618751480198"
        # freq="0.00552425 0.00781246 0.0110485 0.0156249 0.022097 0.03125 0.04419417382415922 0.0625 0.08838834764831845 0.125 0.1767766952966369 0.25 0.3535533905932738 0.5 0.7071067811865476 1.0 1.4142135623730951 2.0 2.8284271247461903 4.0 5.656854249492381 8.0 11.313708498984761 16.0 22.627416997969522 32.0 45.254833995939045 64.0 90.50966799187809 128.0 181.01933598375618 256.0 362.03867196751236 512.0 724.0773439350247 1024.0"
        freq=[float(det) for _,det in enumerate(freq.split())]
        return 1.0/np.array(freq)

class PARAM:
    def __init__(self):
        self.STEP,self.START,self.END=1,1,1
        self.END=self.END+self.STEP
        self.ROOT,self.SETNAME,self.TEMP="datasetsedd","train","temp"

        self.init()

    def init(self):
        if os.path.exists(self.set_path()) is False:
            os.mkdir(self.set_path())

    def set_range(self):
        return range(self.START,self.END,self.STEP)
    
    def set_path(self):
        return os.path.join(self.ROOT,self.SETNAME)
    
    def param_name(self,pid):
        return os.path.join(self.set_path(),pid+".dat")
    
    def forward_name(self,pid):
        return os.path.join(self.set_path(),pid+".res")
    
    def pred_name(self,pid):
        return os.path.join(self.set_path(),pid+".pred")
    
    def temp_measured_name(self):
        return os.path.join(self.ROOT,self.TEMP,self.TEMP+".mea")
    
    def temp_knowed_name(self):
        return os.path.join(self.ROOT,self.TEMP,self.TEMP+".mod")
    
    def temp_resulted_name(self):
        return os.path.join(self.ROOT,self.TEMP,self.TEMP+".res")
    
    def is_train_name(self):
        return self.SETNAME in ["train","tra"]

def INITSTEP(f,p,k,m):
    global _forward,_param,_know,_mode
    _forward,_param,_know,_mode=f,p,k,m

def FIRSTSTEP(indexofpid):
    # zoom,original,target,MAXGROUND,NOISE=1,64,256,1200.0,np.random.randint(low=0,high=6)/100.0
    zoom,original,target,MAXGROUND,NOISE=1,64,256,1200.0,0/100.0
    length,high,ground,value=_mode[indexofpid]
    
    for _ in _param.set_range():
        pid=str(uuid.uuid4())
        # pid="0f3e1215-9875-43f9-a2fe-a273c4d7e101"
        path=_param.set_path()

        length=np.random.randint(low=length[0],high=length[1])
        high=np.random.randint(low=high[0],high=high[1])
        sign=np.random.randint(low=0,high=1)
        if sign==1:
            length,high=high,length

        # coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
        # coory=np.random.randint(low=0+3*high+high/2,high=_know.NLAYER-_know.ZEXPAND-high/2)
        gdv=np.full((_know.NLAYER,_know.XRANGE),1.0-ground/MAXGROUND)
        body=np.full((_know.NLAYER,_know.XRANGE),6)

        param=((8,11,3),(12,16,3),(20,26,1))

        LAYER=param[0]
        index=np.random.randint(low=LAYER[0],high=LAYER[1])
        layer00,layer02=index,index+LAYER[2]
        Fixed=range(0,_know.XRANGE,20)
        tmp=[]
        for i in Fixed:
            tmp.append(float(np.random.randint(low=layer00,high=layer02)))
        f=interp1d(Fixed,tmp,axis=0,fill_value="extrapolate")
        paramed0=[int(i) for i in f(range(0,_know.XRANGE,1))]
        j=0
        for row in paramed0:
            for i in range(0,row,1):
                gdv[i][j]=1.0-value[0]/MAXGROUND
                body[i][j]=0
            j=j+1

        LAYER=param[1]
        index=np.random.randint(low=LAYER[0],high=LAYER[1])
        layer00,layer02=index,index+LAYER[2]
        Fixed=range(0,_know.XRANGE,20)
        tmp=[]
        for i in Fixed:
            tmp.append(float(np.random.randint(low=layer00,high=layer02)))
        f=interp1d(Fixed,tmp,axis=0,fill_value="extrapolate")
        paramed2=[int(i) for i in f(range(0,_know.XRANGE,1))]
        j=0
        for row in paramed2:
            for i in range(paramed0[j],row,1):
                gdv[i][j]=1.0-value[1]/MAXGROUND
                body[i][j]=1
            j=j+1

        LAYER=param[2]
        index=np.random.randint(low=LAYER[0],high=LAYER[1])
        layer00,layer02=index,index+LAYER[2]
        Fixed=range(0,_know.XRANGE,20)
        tmp=[]
        for i in Fixed:
            tmp.append(float(np.random.randint(low=layer00,high=layer02)))
        f=interp1d(Fixed,tmp,axis=0,fill_value="extrapolate")
        paramed3=[int(i) for i in f(range(0,_know.XRANGE,1))]
        j=0
        for row in paramed3:
            for i in range(paramed2[j],row,1):
                gdv[i][j]=1.0-value[2]/MAXGROUND
                body[i][j]=2
            j=j+1

        j=0
        for row in paramed3:
            for i in range(paramed3[j],_know.NLAYER,1):
                gdv[i][j]=1.0-value[3]/MAXGROUND
                body[i][j]=3
            j=j+1
        
        sign=np.random.randint(low=1,high=2)
        res=np.random.uniform(value[4][0],value[4][1])
        if sign==0:
            coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
            coory=np.random.randint(low=paramed0[coorx],high=paramed2[coorx]-high)
            for j in range(length):
                for i in range(high):
                    gdv[i+coory][j+coorx]=1.0-res/MAXGROUND
                    body[i+coory][j+coorx]=4            
        elif sign==1:
            coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
            coory=np.random.randint(low=paramed2[coorx],high=paramed3[coorx])
            for j in range(length):
                for i in range(high):
                    gdv[i+coory][j+coorx]=1.0-res/MAXGROUND
                    body[i+coory][j+coorx]=4
        elif sign==2:
            pass
        elif sign==3:
            coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
            coory=np.random.randint(low=paramed3[coorx],high=_know.NLAYER-_know.ZEXPAND-high)
            for j in range(length):
                for i in range(high):
                    gdv[i+coory][j+coorx]=1.0-res/MAXGROUND
                    body[i+coory][j+coorx]=4

        sign=np.random.randint(low=3,high=4)
        res=np.random.uniform(value[4][0],value[4][1])
        if sign==0:
            coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
            coory=np.random.randint(low=paramed0[coorx],high=paramed2[coorx]-high)
            for j in range(length):
                for i in range(high):
                    gdv[i+coory][j+coorx]=1.0-res/MAXGROUND
                    body[i+coory][j+coorx]=4            
        elif sign==1:
            coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
            coory=np.random.randint(low=paramed2[coorx],high=paramed3[coorx])
            for j in range(length):
                for i in range(high):
                    gdv[i+coory][j+coorx]=1.0-res/MAXGROUND
                    body[i+coory][j+coorx]=4
        elif sign==2:
            pass
        elif sign==3:
            coorx=np.random.randint(low=_know.XEXPAND+length,high=_know.XRANGE-_know.XEXPAND-length)
            coory=np.random.randint(low=paramed3[coorx],high=_know.NLAYER-_know.ZEXPAND-high)
            for j in range(length):
                for i in range(high):
                    gdv[i+coory][j+coorx]=1.0-res/MAXGROUND
                    body[i+coory][j+coorx]=4

        nx,ny,data=_know.XRANGE,_know.NLAYER,[]
        image=np.ones((ny*zoom,nx*zoom),dtype=np.float)
        for j in range(_know.NLAYER):
            data.append([float(i) for i in gdv[j]])
        for i in range(_know.NLAYER):
            for j in range(_know.XRANGE):
                image[i*zoom:(i+1)*zoom,j*zoom:(j+1)*zoom]=data[i][j]
        
        top_border,bottom_border,left_border,right_border=int((target-original)/2.),int((target-original)/2.),int((target-original)/2.),int((target-original)/2.)
        image=cv2.copyMakeBorder(image,int((target-ny*zoom)/2),int((target-ny*zoom)/2),int((target-nx*zoom)/2),int((target-nx*zoom)/2),cv2.BORDER_REPLICATE)

        x0,x2=np.random.randint(low=0,high=int(target/2)),np.random.randint(low=int(target/2),high=target)
        sign,fault,interval=np.random.randint(low=0,high=2),1 if np.random.randint(low=0,high=2)==0 else -1,np.random.randint(low=3,high=7)
        if sign==1:
            x0,x2=x2,x0

        f=interp1d((0,target),(x0,x2))
        swap=np.copy(image)
        for i in range(0+interval,target-interval,1):
            x0,y0=int(f(i)),i
            x2,y2=int(f(i+interval*fault)),i+interval*fault
            if x0<=x2:
                image[y0,x0:(x0+target-x2)]=swap[y2,x2:]
            elif x0>x2:
                image[y0,x0:]=swap[y2,x2:(x2+target-x0)]
                
        image=image[top_border:target-bottom_border,left_border:target-right_border]
        np.savetxt(_param.param_name(pid),image)

        curr_proc=mp.current_process()
        result=_forward.Forward(np.array(image).astype(np.float32),_param.temp_knowed_name(),_param.temp_resulted_name(),original,zoom,nx,ny,_know.XRANGE,_know.NLAYER,_know.INTERVAL,_know.TIC,MAXGROUND,curr_proc._identity[0])

        data=result[:int(len(result)/2)]
        TEREAL=np.array([det for i,det in enumerate(data) if i%2==0]).reshape((len(_know.COOR),_forward.NFREQ))
        # TEREAL=np.random.normal(TEREAL,NOISE*TEREAL.max())
        TEREAL=np.random.normal(TEREAL,NOISE*TEREAL)
        # cv2.imwrite(os.path.join(path,pid+"_tereal.jpg"),TEREAL*255)
        TEIMAG=np.array([det for i,det in enumerate(data) if i%2==1]).reshape((len(_know.COOR),_forward.NFREQ))
        # TEIMAG=np.random.normal(TEIMAG,NOISE*TEIMAG.max())
        TEIMAG=np.random.normal(TEIMAG,NOISE*TEIMAG)
        # cv2.imwrite(os.path.join(path,pid+"_teimag.jpg"),TEIMAG*255)

        content=[]
        TEREAL,TEIMAG=TEREAL.reshape((1,len(_know.COOR)*_forward.NFREQ))[0],TEIMAG.reshape((1,len(_know.COOR)*_forward.NFREQ))[0]
        for i,_ in enumerate(TEREAL):
            content.extend([TEREAL[i],TEIMAG[i]])

        data=result[int(len(result)/2):]
        TMREAL=np.array([det for i,det in enumerate(data) if i%2==0]).reshape((len(_know.COOR),_forward.NFREQ))
        # TMREAL=np.random.normal(TMREAL,NOISE*TMREAL.max())
        TMREAL=np.random.normal(TMREAL,NOISE*TMREAL)
        # cv2.imwrite(os.path.join(path,pid+"_tmreal.jpg"),TMREAL*255)
        TMIMAG=np.array([det for i,det in enumerate(data) if i%2==1]).reshape((len(_know.COOR),_forward.NFREQ))
        # TMIMAG=np.random.normal(TMIMAG,NOISE*TMIMAG.max())
        TMIMAG=np.random.normal(TMIMAG,NOISE*TMIMAG)
        # cv2.imwrite(os.path.join(path,pid+"_tmimag.jpg"),TMIMAG*255)

        TMREAL,TMIMAG=TMREAL.reshape((1,len(_know.COOR)*_forward.NFREQ))[0],TMIMAG.reshape((1,len(_know.COOR)*_forward.NFREQ))[0]
        for i,_ in enumerate(TMREAL):
            content.extend([TMREAL[i],TMIMAG[i]])

        with open(_param.forward_name(pid),"w") as fp:
            fp.write("\n".join([str(det) for det in content]))

        # cv2.imwrite(os.path.join(path,pid+".jpg"),image*255)
        # if os.path.exists(os.path.join(path,pid)) is False:
        #     os.mkdir(os.path.join(path,pid))
        #     os.mkdir(os.path.join(path,pid,"inversion"))
        #     os.mkdir(os.path.join(path,pid,"forward"))
        
        # i,know,result=str(curr_proc._identity[0]),_param.temp_knowed_name(),_param.temp_resulted_name()
        # shutil.copyfile(know+i,os.path.join(path,pid,"forward",os.path.basename(know)))
        # shutil.copyfile(result+i,os.path.join(path,pid,"forward",os.path.basename(result)))

        # files=["mt2di.exe","temp.mod","control.inv"]
        # for each in files:
        #     shutil.copyfile(each,os.path.join(path,pid,"inversion",each))

        # content=[]
        # with open(os.path.join(path,pid,"forward",os.path.basename(result)),"r") as fp:
        #     content.extend([det.split() for det in fp])
        
        # swap=content[:]

        # content=swap[:int(len(swap)/2)]

        # header=content[:8]
        # data=df(content[8:],columns=header[1][1:])
        # data["Real"]=data["Real"].astype(float)
        # data["Imag"]=data["Imag"].astype(float)

        # TEREAL=np.array(data["Real"].values).reshape((len(_know.COOR),_forward.NFREQ))
        # TEREAL=np.random.normal(TEREAL,NOISE*abs(TEREAL.max()))
        # data["Real"]=TEREAL.reshape((1,len(_know.COOR)*_forward.NFREQ))[0]

        # TEIMAG=np.array(data["Imag"].values).reshape((len(_know.COOR),_forward.NFREQ))
        # TEIMAG=np.random.normal(TEIMAG,NOISE*abs(TEIMAG.max()))
        # data["Imag"]=TEIMAG.reshape((1,len(_know.COOR)*_forward.NFREQ))[0]

        # noise=NOISE
        # if NOISE==0.0:
        #     NOISE=0.05

        # data["Error"]=data.apply(lambda x: sqrt(x["Real"]*x["Real"]+x["Imag"]*x["Imag"])*NOISE,axis=1)
        # NOISE=noise

        # data["Real"]=data["Real"].astype(str)
        # data["Imag"]=data["Imag"].astype(str)
        # data["Error"]=data["Error"].astype(str)

        # with open(os.path.join(path,pid,"inversion","temp.res"),"w") as fp:
        #     fp.write("\n".join([" ".join(each) for each in header])+"\n")
        #     fp.write("\n".join([" ".join(each) for each in data.values])+"\n")

        # content=swap[int(len(swap)/2):]

        # header=content[:8]
        # data=df(content[8:],columns=header[1][1:])
        # data["Real"]=data["Real"].astype(float)
        # data["Imag"]=data["Imag"].astype(float)
        # TMREAL=np.array(data["Real"].values).reshape((len(_know.COOR),_forward.NFREQ))
        # TMREAL=np.random.normal(TMREAL,NOISE*abs(TMREAL.max()))
        # data["Real"]=TMREAL.reshape((1,len(_know.COOR)*_forward.NFREQ))[0]

        # TMIMAG=np.array(data["Imag"].values).reshape((len(_know.COOR),_forward.NFREQ))
        # TMIMAG=np.random.normal(TMIMAG,NOISE*abs(TMIMAG.max()))
        # data["Imag"]=TMIMAG.reshape((1,len(_know.COOR)*_forward.NFREQ))[0]

        # noise=NOISE
        # if NOISE==0.0:
        #     NOISE=0.05

        # data["Error"]=data.apply(lambda x: sqrt(x["Real"]*x["Real"]+x["Imag"]*x["Imag"])*NOISE,axis=1)     
        # NOISE=noise   

        # data["Real"]=data["Real"].astype(str)
        # data["Imag"]=data["Imag"].astype(str)
        # data["Error"]=data["Error"].astype(str)

        # with open(os.path.join(path,pid,"inversion","temp.res"),"a") as fp:
        #     fp.write("\n".join([" ".join(each) for each in header])+"\n")
        #     fp.write("\n".join([" ".join(each) for each in data.values])+"\n")
    
if __name__=='__main__':
    param=PARAM()
    measured=MEASURE()
    know=KNOW(36,50.0,1.0,10,50,150.0,15.0,5.0)

    measured.Measured(param.temp_measured_name(),know.COOR)

    mode=([(10,16),(4,7),1000.0,(350.,450.,600.,800.,(1.0,6.0),30.0,70.0)],)*4000
    # mode=([(6,10),(4,7),1000.0,(190.,400.,600.,800.,(1.0,20.0),30.0,70.0)],)*4000
    # mode=([(10,16),(4,7),100.0,(100.,100.,100.,100.,(100.,101.),30.0,70.0)],)*1

    #mp.cpu_count()
    pool=mp.Pool(16,initializer=INITSTEP,initargs=(measured,param,know,mode))
    pool.map(FIRSTSTEP,[i for i in range(len(mode))])
    pool.close()